// Copyright 2021 The Kubeflow Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// Launcher command for Kubeflow Pipelines v2.
package main

import (
	"context"
	"flag"
	"fmt"

	"github.com/golang/glog"
	"github.com/kubeflow/pipelines/backend/src/v2/client_manager"
	"github.com/kubeflow/pipelines/backend/src/v2/component"
	"github.com/kubeflow/pipelines/backend/src/v2/config"
)

// TODO: use https://github.com/spf13/cobra as a framework to create more complex CLI tools with subcommands.
var (
	copy                    = flag.String("copy", "", "copy this binary to specified destination path")
	pipelineName            = flag.String("pipeline_name", "", "pipeline context name")
	runID                   = flag.String("run_id", "", "pipeline run uid")
	parentDagID             = flag.Int64("parent_dag_id", 0, "parent DAG execution ID")
	executorType            = flag.String("executor_type", "container", "The type of the ExecutorSpec")
	executionID             = flag.Int64("execution_id", 0, "Execution ID of this task.")
	executorInputJSON       = flag.String("executor_input", "", "The JSON-encoded ExecutorInput.")
	componentSpecJSON       = flag.String("component_spec", "", "The JSON-encoded ComponentSpec.")
	importerSpecJSON        = flag.String("importer_spec", "", "The JSON-encoded ImporterSpec.")
	taskSpecJSON            = flag.String("task_spec", "", "The JSON-encoded TaskSpec.")
	podName                 = flag.String("pod_name", "", "Kubernetes Pod name.")
	podUID                  = flag.String("pod_uid", "", "Kubernetes Pod UID.")
	mlPipelineServerAddress = flag.String("ml_pipeline_server_address", "ml-pipeline.kubeflow", "The name of the ML pipeline API server address.")
	mlPipelineServerPort    = flag.String("ml_pipeline_server_port", "8887", "The port of the ML pipeline API server.")
	mlmdServerAddress       = flag.String("mlmd_server_address", "", "The MLMD gRPC server address.")
	mlmdServerPort          = flag.String("mlmd_server_port", "8080", "The MLMD gRPC server port.")
	logLevel                = flag.String("log_level", "1", "The verbosity level to log.")
	publishLogs             = flag.String("publish_logs", "true", "Whether to publish component logs to the object store")
	cacheDisabledFlag       = flag.Bool("cache_disabled", false, "Disable cache globally.")
	caCertPath              = flag.String("ca_cert_path", "", "The path to the CA certificate to trust on connections to the ML pipeline API server and metadata server.")
	mlPipelineTLSEnabled    = flag.Bool("ml_pipeline_tls_enabled", false, "Set to true if mlpipeline API server serves over TLS.")
	metadataTLSEnabled      = flag.Bool("metadata_tls_enabled", false, "Set to true if MLMD serves over TLS.")
)

func main() {
	err := run()
	if err != nil {
		glog.Exit(err)
	}
}

func run() error {
	flag.Parse()
	ctx := context.Background()

	glog.Infof("Setting log level to: '%s'", *logLevel)
	err := flag.Set("v", *logLevel)
	if err != nil {
		glog.Warningf("Failed to set log level: %s", err.Error())
	}

	if *copy != "" {
		// copy is used to copy this binary to a shared volume
		// this is a special command, ignore all other flags by returning
		// early
		return component.CopyThisBinary(*copy)
	}
	namespace, err := config.InPodNamespace()
	if err != nil {
		return err
	}

	launcherV2Opts := &component.LauncherV2Options{
		Namespace:               namespace,
		PodName:                 *podName,
		PodUID:                  *podUID,
		MLPipelineServerAddress: *mlPipelineServerAddress,
		MLPipelineServerPort:    *mlPipelineServerPort,
		MLMDServerAddress:       *mlmdServerAddress,
		MLMDServerPort:          *mlmdServerPort,
		PipelineName:            *pipelineName,
		RunID:                   *runID,
		PublishLogs:             *publishLogs,
		CacheDisabled:           *cacheDisabledFlag,
		MLPipelineTLSEnabled:    *mlPipelineTLSEnabled,
		MLMDTLSEnabled:          *metadataTLSEnabled,
		CaCertPath:              *caCertPath,
	}

	switch *executorType {
	case "importer":
		importerLauncherOpts := &component.ImporterLauncherOptions{
			PipelineName: *pipelineName,
			RunID:        *runID,
			ParentDagID:  *parentDagID,
		}
		importerLauncher, err := component.NewImporterLauncher(ctx, *componentSpecJSON, *importerSpecJSON, *taskSpecJSON, launcherV2Opts, importerLauncherOpts)
		if err != nil {
			return err
		}
		if err := importerLauncher.Execute(ctx); err != nil {
			return err
		}
		return nil
	case "container":
		clientOptions := &client_manager.Options{
			MLPipelineServerAddress: launcherV2Opts.MLPipelineServerAddress,
			MLPipelineServerPort:    launcherV2Opts.MLPipelineServerPort,
			MLMDServerAddress:       launcherV2Opts.MLMDServerAddress,
			MLMDServerPort:          launcherV2Opts.MLMDServerPort,
			CacheDisabled:           launcherV2Opts.CacheDisabled,
			MLMDTLSEnabled:          launcherV2Opts.MLMDTLSEnabled,
			CaCertPath:              launcherV2Opts.CaCertPath,
		}
		clientManager, err := client_manager.NewClientManager(clientOptions)
		if err != nil {
			return err
		}
		launcher, err := component.NewLauncherV2(ctx, *executionID, *executorInputJSON, *componentSpecJSON, flag.Args(), launcherV2Opts, clientManager)
		if err != nil {
			return err
		}
		glog.V(5).Info(launcher.Info())
		if err := launcher.Execute(ctx); err != nil {
			return err
		}

		return nil

	}
	return fmt.Errorf("unsupported executor type %s", *executorType)

}

// Use WARNING default logging level to facilitate troubleshooting.
func init() {
	flag.Set("logtostderr", "true")
	// Change the WARNING to INFO level for debugging.
	flag.Set("stderrthreshold", "WARNING")
}
